from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import requests
import json
from langchain.prompts import PromptTemplate
from starlette.middleware.cors import CORSMiddleware

API_URL = 'http://10.1.161.53:11434/api/chat'
MODEL_NAME = "deepseek-r1:70b"

app = FastAPI()

# 允许所有来源的跨域请求
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有HTTP方法
    allow_headers=["*"]  # 允许所有请求头
)

@app.post("/analyse_sql")
async def analyse_sql(request: Request):
    # 获取请求参数
    data = await request.json()
    question = data.get('question')
    sql_result = data.get('sql_result')

    """基于SQL查询结果生成数据分析报告
    Args:
        question: 需要分析的业务问题
        sql_result: SQL查询结果字符串
        llm: 已初始化的语言模型实例
    Returns:
        包含分析、见解和建议的文本报告
    """
    # 创建分析模板
    analysis_prompt = PromptTemplate.from_template(
        """根据问题{test_question}，SQL查询结果{sql_result}，  
        请提供以下内容：  
        1. 对数据的简明分析  
        2. 从数据中得出的关键见解  
        3. 根据数据可能的建议或行动计划  

        请使用简洁专业的语言回答。"""
    )

    # 格式化PromptTemplate为字符串
    formatted_prompt = analysis_prompt.format(
        test_question=question,
        sql_result=sql_result
    )

    # 构建API请求数据
    api_data = {
        "model": MODEL_NAME,
        "messages": [
            {"role": "system", "content": "你是一个专业的数据分析助理，擅长解读SQL查询结果并提出业务建议"},
            {"role": "user", "content": formatted_prompt}
        ],
        "stream": True
    }

    # 发送流式请求
    response = requests.post(
        API_URL,
        data=json.dumps(api_data),
        stream=True,
        headers={'Content-Type': 'application/json'}
    )

    # 流式响应生成器
    def generate():
        try:
            for chunk in response.iter_lines():
                if chunk:
                    decoded = json.loads(chunk.decode('utf-8'))
                    if 'message' in decoded and 'content' in decoded['message']:
                        content = decoded['message']['content']
                        yield f"data: {json.dumps({'content': content})}\n\n"
            yield "data: [DONE]\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=3001)